import base64
import traceback

from ..common import VolcengineClient
from ..util import convert_bytes_to_image, convert_image_to_bytes


class PrettyFaceSwap:
    RETURN_TYPES = ("IMAGE",)
    # RETURN_NAMES = ("image_output_name",)

    FUNCTION = "image_creator"

    # OUTPUT_NODE = False

    CATEGORY = "PrettyFaceSwap"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "info": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "{}",
                    "lazy": True
                }),
                "image": ("IMAGE",),
                "beauty_level": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}),
            },
        }

    @staticmethod
    def image_creator(image, info, beauty_level):

        try:
            client = VolcengineClient(volcengine_info=info)
            if isinstance(image, list):
                images = []
                for img in image:
                    image_bytes = convert_image_to_bytes(img)
                    img_base64 = base64.b64encode(image_bytes).decode("utf-8")
                    operate_result = client.pretty_to_face(image_base64=img_base64,
                                                           beauty_level=beauty_level)
                    if operate_result.code == 1:
                        pics = operate_result.data
                        image_bytes = pics[0]
                        pic = convert_bytes_to_image(image_bytes)
                        images.append(pic)
                image = images
            else:
                image_bytes = convert_image_to_bytes(image)
                image_base64 = base64.b64encode(image_bytes).decode("utf-8")
                operate_result = client.pretty_to_face(image_base64=image_base64,
                                                       beauty_level=beauty_level)
                if operate_result.code == 1:
                    pics = operate_result.data
                    image_bytes = pics[0]
                    image = convert_bytes_to_image(image_bytes)
        except Exception as e:
            print(e, traceback.format_exc())
        return (image,)
